import numpy as np
import cv2
import json
import os


def calculateSquareBoundingBox(
    points: np.ndarray,
    originalWidth: int,
    originalHeight: int,
    scale: float = 1.5,
) -> np.ndarray:
    """
    根据一张图中多个点的坐标points，计算出一个正方形的边界框，并返回该边界框的四个顶点坐标。

    :param points: 点坐标，shape=(n, 2)
    :param originalWidth: 图像的宽度
    :param originalHeight: 图像的高度
    :param scale: 实际裁切正方形边长为外接正方形边长的scale倍
    :return:
    """

    # 获取外界矩形的e
    rectangle_x_min = round(min(points[:, 0]))
    rectangle_x_max = round(max(points[:, 0]))
    rectangle_y_min = round(min(points[:, 1]))
    rectangle_y_max = round(max(points[:, 1]))

    rectangleWidth = rectangle_x_max - rectangle_x_min
    rectangleHeight = rectangle_y_max - rectangle_y_min

    # 计算正方形边长,并扩充为scale倍
    squareWidth = round(max(rectangleWidth, rectangleHeight) * scale)

    horizontalPadding = squareWidth - rectangleWidth
    verticalPadding = squareWidth - rectangleHeight

    leftPadding = horizontalPadding // 2.0

    rightPadding = horizontalPadding - leftPadding

    topPadding = verticalPadding // 2.0

    bottomPadding = verticalPadding - topPadding

    if rectangle_x_min - leftPadding < 0:
        leftPadding = rectangle_x_min - 0.0
        rightPadding = horizontalPadding - leftPadding
    elif rectangle_x_max + rightPadding > originalWidth:
        rightPadding = originalWidth - rectangle_x_max
        leftPadding = horizontalPadding - rightPadding
    if rectangle_y_min - topPadding < 0:
        topPadding = rectangle_y_min - 0.0
        bottomPadding = verticalPadding - topPadding
    elif rectangle_y_max + bottomPadding > originalHeight:
        bottomPadding = originalHeight - rectangle_y_max
        topPadding = verticalPadding - bottomPadding

    square_x_min = round(np.maximum((rectangle_x_min - leftPadding), 0))
    square_x_max = round(np.minimum((rectangle_x_max + rightPadding), originalWidth))
    square_y_min = round(np.maximum((rectangle_y_min - topPadding), 0))
    square_y_max = round(np.minimum((rectangle_y_max + bottomPadding), originalHeight))
    return np.array(
        [
            [square_x_min, square_y_min],
            [square_x_max, square_y_min],
            [square_x_max, square_y_max],
            [square_x_min, square_y_max],
        ],
        dtype=np.int32,
    )


if __name__ == "__main__":
    # 带有labelme格式的标签和图像的数据集路径
    source_path = "x1_shoted_new"

    # 输出的条码图像数据集路径
    obj_path = "real_barcode_data"
    os.makedirs(obj_path, exist_ok=True)
    name_list = os.listdir(source_path)
    for name in name_list:
        if name.split(".")[-1] == "json":
            with open(os.path.join(source_path, name), "r") as f:
                label_data = json.load(f)

            # 根据json文件中的信息，读取图片数据
            image = cv2.imread(os.path.join(source_path, label_data["imagePath"]))
            image_height = label_data["imageHeight"]
            image_width = label_data["imageWidth"]

            # 获取关键点坐标
            shapes = label_data["shapes"]

            # 遍历所有条码，计算正方形边界框，并保存条码图像和关键点信息
            for shape in shapes:
                if shape["label"].upper() in [
                    "CODE128",
                    "CODE39",
                    "DM",
                    "QR",
                    "EAN13_GUARD",
                    "EAN8_GUARD",
                ]:
                    # 根据关键点信息计算正方形边界框
                    points = np.array(shape["points"])

                    # 对于透视变换后有部分落在图像外面的条码，不进行裁切，直接跳过
                    if (
                        np.any(points < 0)
                        or np.any(points[:, 0] >= image_width)
                        or np.any(points[:, 1] >= image_height)
                    ):
                        continue
                    else:
                        square_points = calculateSquareBoundingBox(
                            points, image_width, image_height
                        )
                        # 创建labelme格式的字典，用于填入正方形裁切图像中的条码关键点信息
                        barcode_image_width = int(
                            square_points[1][0] - square_points[0][0]
                        )
                        barcode_image_height = int(
                            square_points[2][1] - square_points[0][1]
                        )
                        if barcode_image_width != barcode_image_height:
                            print(label_data["imagePath"])
                        exist_data_num = len(os.listdir(obj_path)) // 2
                        label_dict = {
                            "imagePath": str(exist_data_num) + ".png",
                            "imageWidth": barcode_image_width,
                            "imageHeight": barcode_image_height,
                            "imageData": None,
                            "shapes": [],
                        }
                        # 遍历所有条码，若条码有一个关键点落在其对应的正方形边界框内，则添加条码的关键点信息
                        for shape_inner in shapes:
                            if shape_inner["label"].upper() in [
                                "CODE128",
                                "CODE39",
                                "DM",
                                "QR",
                                "EAN13_GUARD",
                                "EAN8_GUARD",
                            ]:
                                # 将条码在原图中的坐标变换到在正方形边界框中的坐标
                                barcode_points = (
                                    np.array(shape_inner["points"])
                                    - square_points[0]
                                    + 1
                                ).astype(np.int32)

                                # 遍历条码，判定是否落在正方形边界框内，若落在，则添加条码的关键点信息
                                for barcode_point in barcode_points:
                                    if (
                                        0 < barcode_point[0] < barcode_image_width
                                        and 0 < barcode_point[1] < barcode_image_height
                                    ):
                                        # 添加条码的关键点信息并保存为字典格式
                                        label_dict["shapes"].append(
                                            {
                                                "label": shape_inner["label"],
                                                "points": barcode_points.tolist(),
                                            }
                                        )
                                        break
                                with open(
                                    os.path.join(
                                        obj_path, str(exist_data_num) + ".json"
                                    ),
                                    "w",
                                ) as f:
                                    json.dump(label_dict, f)
                                # 从原图中裁切出条码图像并保存
                                barcode_image = image[
                                    square_points[0][1] : square_points[2][1],
                                    square_points[0][0] : square_points[1][0],
                                ]
                                cv2.imwrite(
                                    os.path.join(
                                        obj_path, str(exist_data_num) + ".png"
                                    ),
                                    barcode_image,
                                )
